import os
import os.path as osp
import argparse
import onnx
import onnxruntime as rt
import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
from utils import get_test_dataset


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default=osp.join('data', 'mnist_net.onnx'))
    parser.add_argument('--dataset', type=str, choices=['mnist', 'cifar10'])
    args = parser.parse_args()

    sess = rt.InferenceSession(args.model)
    input_name = sess.get_inputs()[0].name
    label_name = sess.get_outputs()[0].name

    _, test_loader = get_test_dataset(args.dataset)

    for x, label in test_loader:
        image = np.ascontiguousarray(x.numpy())
        pred = sess.run([label_name], {input_name: image.astype(np.float32)})[0][0]
        print('label={}, y={}'.format(label, pred.argmax()))
